from com.bosssoft.handler.SentenceHandler import SentenceHandler

class SimilarityUtil(object):
    def __init__(self, modelInstance):
        self.modelInstance = modelInstance

    def sentence2vec(self, sentence):
        sentence = SentenceHandler(sentence, self)
        vec_bow = self.modelInstance.dictionary.doc2bow(sentence.get_cut_sentence())
        return self.modelInstance.model[vec_bow]

    # 求与目标句子最相似的句子
    def similarity(self, sentence):
        # 将目标句子转换为向量
        sentence_vec = self.sentence2vec(sentence)

        sims = self.modelInstance.index[sentence_vec]
        # 求最相似的句子
        sim = max(enumerate(sims), key=lambda item: item[1])

        index = sim[0]
        score = sim[1]
        sentence = self.modelInstance.sentences[index]

        sentence.set_score(score)
        return sentence

    # 求与目标句子TopN相似的句子
    def topNSimilarity(self, sentence, n):
        # 将目标句子转换为向量
        sentence_vec = self.sentence2vec(sentence)
        sims = self.modelInstance.index.get_similarities(sentence_vec)

        # 我们根据相似度进行排序，并取前 N 个
        sorted_sims = sorted(enumerate(sims), key=lambda item: item[1], reverse=True)

        # 提取前 N 个最相似的句子
        top_n_sentences = [self.modelInstance.sentences[index] for index, _ in sorted_sims[:n]]

        # 设置每个句子的相似度分数
        for i, sentence in enumerate(top_n_sentences):
           sentence.set_score(sorted_sims[i][1])

        return top_n_sentences